#include <stdio.h>
#include <fstream>

#include "../params.h"
#include "../GridException.h"
#include "../qmc.h"
#include "clUtil/clUtil.h"

#define PARTICLE_MASS 7296.29935
#define SHOW_DEVICE

using namespace std;

//Storage needed for computing wavefunction, potential, and kinetic
static float* partialPotential;
static float* partialWaveFunction;
static double* partialKinetic;
static double* qForceX;
static double* qForceY;
static double* qForceZ;
static cl_mem xGPU;
static cl_mem yGPU;
static cl_mem zGPU;
static cl_mem partialKineticGPU;
static cl_mem qForceXGPU;
static cl_mem qForceYGPU;
static cl_mem qForceZGPU;
static cl_mem partialWaveFunctionGPU;
static cl_mem partialPotentialGPU;
  

//Perform reduction on a vector of floats
static float reduce(float* a, int n)
{
  double total = 0.0;

  for(int i = 0; i < n; i++)
  {
    total += (double)a[i];
  }

  return (float)total;
}

//Perform reduction on a vector of doubles
static float reduce(double* a, int n)
{
  double total = 0.0;

  for(int i = 0; i < n; i++)
  {
    total += (double)a[i];
  }

  return (float)total;
}


bool GridOpenCLInit()
{
  const char* kernelNames[] = {"OpenCL/kernel.cl"};
  char options[1024];
  char* mchamrPath;
  char const* waveFunctionName;
  size_t numParticles = gXDim * gYDim * gZDim;

  //Get the path to mchamr
  mchamrPath = getenv("MCHAMR_PATH");
  
  //chdir into mchamr directory
  chdir(mchamrPath);

  //create header file containing constants
  generateParams();

  //Get which wavefunction we're using
  switch(getWaveFunction())
  {
    case kWaveFunctionDefault:
      waveFunctionName = "DEFAULT_WAVE_FUNCTION";
      break;
    case kWaveFunctionMcMillan:
      waveFunctionName = "MCMILLAN_WAVE_FUNCTION";
      break;
    default:
      printf("Error GridOpenCLInit(): Invalid wave function.\n");
      break;
  }

  snprintf(options, 
           sizeof(options) - 1, "-D%s -I%s",
           waveFunctionName,
           mchamrPath);

  if(clUtil::initialize(kernelNames, 1, NULL, options) != CL_SUCCESS)
  {
    printf("Could not initialize OpenCL runtime. Aborting.");
    exit(1);
  }
  
  //Allocate memory for calculations
  partialPotential = new float[numParticles];
  partialWaveFunction = new float[numParticles];
  partialKinetic = new double[numParticles];
  qForceX = new double[numParticles];
  qForceY = new double[numParticles];
  qForceZ = new double[numParticles];

  clUtil::alloc(numParticles * sizeof(partialPotential[0]), 
                &partialPotentialGPU);
  clUtil::alloc(numParticles * sizeof(partialWaveFunction[0]), 
                &partialWaveFunctionGPU);
  clUtil::alloc(numParticles * sizeof(float), &xGPU);
  clUtil::alloc(numParticles * sizeof(float), &yGPU);
  clUtil::alloc(numParticles * sizeof(float), &zGPU);
  clUtil::alloc(numParticles * sizeof(partialKinetic[0]), &partialKineticGPU);
  clUtil::alloc(numParticles * sizeof(qForceX[0]), &qForceXGPU);
  clUtil::alloc(numParticles * sizeof(qForceY[0]), &qForceYGPU);
  clUtil::alloc(numParticles * sizeof(qForceZ[0]), &qForceZGPU);

}

bool GridOpenCLDeinit()
{
  delete[] partialPotential;
  delete[] partialWaveFunction;
  delete[] partialKinetic;
  delete[] qForceX;
  delete[] qForceY;
  delete[] qForceZ;

  clUtil::free(xGPU);
  clUtil::free(yGPU);
  clUtil::free(zGPU);
  clUtil::free(partialPotentialGPU);
  clUtil::free(partialWaveFunctionGPU);
  clUtil::free(xGPU);
  clUtil::free(yGPU);
  clUtil::free(zGPU);
  clUtil::free(partialKineticGPU);
  clUtil::free(qForceXGPU);
  clUtil::free(qForceYGPU);
  clUtil::free(qForceZGPU);

}

static void computeKineticEnergy(int numParticles,
                                 float* x,
                                 float* y,
                                 float* z,
                                 float* kinetic)
{
  static cl_kernel kernel = 0;
  cl_int err;
  size_t localWork = 64;
  size_t globalWork = 
    numParticles % localWork == 0 && numParticles > localWork ? 
    numParticles :
    (numParticles / localWork + 1) * localWork;
  double kineticAccumulator = 0.0;


  //Get the kernel we need to compute the wave function
  if(kernel == 0)
  {
    kernel = clCreateKernel(clUtil::getProgram(), "computeKineticEnergy", &err);
    clUtilCheckErrorVoid(err);
  }

  //Copy coordinates to GPU
  clUtil::devicePut(x, numParticles * sizeof(x[0]), xGPU);
  clUtil::devicePut(y, numParticles * sizeof(y[0]), yGPU);
  clUtil::devicePut(z, numParticles * sizeof(z[0]), zGPU);

  err = clSetKernelArg(kernel, 0, sizeof(xGPU), &xGPU);
  clUtilCheckErrorVoid(err);
  err = clSetKernelArg(kernel, 1, sizeof(yGPU), &yGPU);
  clUtilCheckErrorVoid(err);
  err = clSetKernelArg(kernel, 2, sizeof(zGPU), &zGPU);
  clUtilCheckErrorVoid(err);
  err = clSetKernelArg(kernel, 3, sizeof(numParticles), &numParticles);
  clUtilCheckErrorVoid(err);
  err = clSetKernelArg(kernel, 4, sizeof(qForceXGPU), &qForceXGPU);
  clUtilCheckErrorVoid(err);
  err = clSetKernelArg(kernel, 5, sizeof(qForceYGPU), &qForceYGPU);
  clUtilCheckErrorVoid(err);
  err = clSetKernelArg(kernel, 6, sizeof(qForceZGPU), &qForceZGPU);
  clUtilCheckErrorVoid(err);
  err = clSetKernelArg(kernel, 
                       7, 
                       sizeof(partialKineticGPU), 
                       &partialKineticGPU);
  clUtilCheckErrorVoid(err);

  err = clEnqueueNDRangeKernel(clUtil::getCommandQueue(),
                               kernel,
                               1,
                               NULL,
                               &globalWork,
                               &localWork,
                               0,
                               NULL,
                               NULL);
  clUtilCheckErrorVoid(err);

  //Fetch computed partial results
  clUtil::deviceGet(partialKinetic, 
                    numParticles * sizeof(partialKinetic[0]),
                    partialKineticGPU);
  clUtilCheckErrorVoid(err);
  
  clUtil::deviceGet(qForceX, 
                    numParticles * sizeof(qForceX[0]),
                    qForceXGPU);
  clUtilCheckErrorVoid(err);
  
  clUtil::deviceGet(qForceY, 
                    numParticles * sizeof(qForceY[0]),
                    qForceYGPU);
  clUtilCheckErrorVoid(err);
  
  clUtil::deviceGet(qForceZ, 
                    numParticles * sizeof(qForceZ[0]),
                    qForceZGPU);
  clUtilCheckErrorVoid(err);

  //Sum kinetic contributions from each particle
  kineticAccumulator = reduce(partialKinetic, numParticles);
  kineticAccumulator *= 2.0;

  //Subract q-force
  for(int i = 0; i < numParticles; i++)
  {
    kineticAccumulator -= qForceX[i] * qForceX[i] +
                          qForceY[i] * qForceY[i] +
                          qForceZ[i] * qForceZ[i];
  } 

  kineticAccumulator *= 0.5 / PARTICLE_MASS;
  *kinetic = (float)kineticAccumulator;
}

static void computeWaveFunction(int numParticles,
                                float* x,
                                float* y,
                                float* z,
                                float* potential,
                                float* waveFunction)
{
  static cl_kernel kernel = 0;
  cl_int err;
  size_t localWork = 64;
  size_t globalWork = 
    numParticles % localWork == 0 ? 
    numParticles :
    (numParticles / localWork + 1) * localWork;


  //Get the kernel we need to compute the wave function
  if(kernel == 0)
  {
    kernel = clCreateKernel(clUtil::getProgram(), "computeWaveFunction", &err);
    clUtilCheckErrorVoid(err);
  }

  //Allocate
  //Copy to GPU
  clUtil::devicePut(x, numParticles * sizeof(float), xGPU);
  clUtil::devicePut(y, numParticles * sizeof(float), yGPU);
  clUtil::devicePut(z, numParticles * sizeof(float), zGPU);

  err = clSetKernelArg(kernel, 0, sizeof(xGPU), &xGPU);
  clUtilCheckErrorVoid(err);
  err = clSetKernelArg(kernel, 1, sizeof(yGPU), &yGPU);
  clUtilCheckErrorVoid(err);
  err = clSetKernelArg(kernel, 2, sizeof(zGPU), &zGPU);
  clUtilCheckErrorVoid(err);
  err = clSetKernelArg(kernel, 3, sizeof(numParticles), &numParticles);
  clUtilCheckErrorVoid(err);
  err = clSetKernelArg(kernel, 
                       4, 
                       sizeof(partialWaveFunctionGPU), 
                       &partialWaveFunctionGPU);
  clUtilCheckErrorVoid(err);
  err = clSetKernelArg(kernel, 
                       5, 
                       sizeof(partialPotentialGPU), 
                       &partialPotentialGPU);
  clUtilCheckErrorVoid(err);

  err = clEnqueueNDRangeKernel(clUtil::getCommandQueue(),
                               kernel,
                               1,
                               NULL,
                               &globalWork,
                               &localWork,
                               0,
                               NULL,
                               NULL);
  clUtilCheckErrorVoid(err);

  clUtil::deviceGet(partialPotential, 
                    numParticles * sizeof(float),
                    partialPotentialGPU);
  clUtilCheckErrorVoid(err);

  clUtil::deviceGet(partialWaveFunction,
                    numParticles * sizeof(float),
                    partialWaveFunctionGPU);
  clUtilCheckErrorVoid(err);
  
  *waveFunction = reduce(partialWaveFunction, numParticles);
  *potential = reduce(partialPotential, numParticles);
}

void GridOpenCLComputeTotalEnergy(int devNum,
                                  float* x,
                                  float* y,
                                  float* z,
                                  int numParticles,
                                  float* waveFunction,
                                  float* potential,
                                  float* kinetic)
{

  computeWaveFunction(numParticles, x, y, z, potential, waveFunction);
  computeKineticEnergy(numParticles, x, y, z, kinetic);
}
